import torch
from FocalLoss.focal_loss_pytorch import FocalLoss
from FocalLoss.focal_loss_mmlab import py_sigmoid_focal_loss

if __name__ == '__main__':
    torch.manual_seed(2020)
    criteria1 = FocalLoss(gamma=2, reduction='mean')
    criteria2 = torch.nn.CrossEntropyLoss()

    x = torch.randn((5, 3), requires_grad=True)
    t = torch.tensor([2, 1, 2, 0, 2], dtype=torch.long)
    print('output: \n', x)
    print('target: \n', t)

    loss1 = criteria1(x, t)
    loss2 = criteria2(x, t)

    loss1.backward()
    loss2.backward()
    # mmdetection's focal loss implemention
    t_one_hot = torch.zeros_like(x).scatter_(1, t.unsqueeze(1), 1)
    loss3 = py_sigmoid_focal_loss(x, t_one_hot, reduction='mean')
    print('my focal loss:', loss1)
    print('cross entropy:', loss2)
    print('mm focal loss:', loss3)

    """
    print below:
    >>> my focal loss: tensor(0.1916, grad_fn=<BinaryCrossEntropyBackward>)
    >>> cross entropy: tensor(1.5936, grad_fn=<NllLossBackward>)
    >>> mm focal loss: tensor(0.1916, grad_fn=<MeanBackward0>)
    
    from the result, we can see, my focal loss implemention is equal to mm's focal loss implemention.
    so enjoy it!
    """